import asyncio
import aiohttp
import poplib
import imaplib
import base64
import re
from email.parser import BytesParser
from email.policy import default
import logging

logger = logging.getLogger(__name__)

# 配置日志（根据需要调整）
logging.basicConfig(level=logging.INFO)

# POP3服务器配置
pop3_server = 'outlook.office365.com'
pop3_port = 995

# 使用 OAuth2 生成认证字符串
def generate_auth_string(user, token):
    auth_string = f"user={user}\1auth=Bearer {token}\1\1"
    return auth_string

# 从正文中提取验证码
def extract_verification_code(body):
    # 根据实际验证码格式调整正则表达式
    code_pattern = r'\b\d{6}\b'  # 例如，提取6位数字的验证码
    match = re.search(code_pattern, body)
    if match:
        return match.group(0).strip()
    else:
        return None

# 解析邮件正文（假设是纯文本或 HTML）
def parse_email_body(msg):
    if msg.is_multipart():
        for part in msg.walk():
            content_type = part.get_content_type()
            if content_type in ['text/plain', 'text/html']:
                charset = part.get_content_charset() or 'utf-8'
                return part.get_payload(decode=True).decode(charset, errors='replace')
    else:
        content_type = msg.get_content_type()
        if content_type in ['text/plain', 'text/html']:
            charset = msg.get_content_charset() or 'utf-8'
            return msg.get_payload(decode=True).decode(charset, errors='replace')
    return None

# 异步获取 access_token
async def get_access_token(client_id, refresh_token):
    url = 'https://login.live.com/oauth20_token.srf'
    data = {
        'client_id': client_id,
        'grant_type': 'refresh_token',
        'refresh_token': refresh_token
    }
    async with aiohttp.ClientSession() as session:
        async with session.post(url, data=data) as response:
            if response.status != 200:
                logger.error(f"获取 access_token 失败，状态码: {response.status}")
                raise Exception("Failed to get access token")
            json_response = await response.json()
            return json_response['access_token']

async def connect_pop3_async(email, client_id, refresh_token, target_sender, retries=6, sleep_time=2):
    loop = asyncio.get_event_loop()

    for attempt in range(1, retries + 1):
        server = None  # 初始化 server 变量
        try:
            logger.info(f"第 {attempt} 次尝试连接 POP3 服务器并获取验证码。")

            # 获取 access_token
            access_token = await get_access_token(client_id, refresh_token)

            # 使用 SSL 连接到 POP3 服务器
            server = await loop.run_in_executor(None, poplib.POP3_SSL, pop3_server, pop3_port)

            # 使用 OAuth2 验证
            auth_string = generate_auth_string(email, access_token)
            encoded_auth_string = base64.b64encode(auth_string.encode("utf-8")).decode("utf-8")

            await loop.run_in_executor(None, server._shortcmd, 'AUTH XOAUTH2')
            await loop.run_in_executor(None, server._shortcmd, encoded_auth_string)

            # 获取邮件列表
            response, listings, octets = await loop.run_in_executor(None, server.list)
            num_messages = len(listings)
            logger.info(f"收件箱中共有 {num_messages} 封邮件。")

            # 从最新的邮件开始遍历
            for i in range(num_messages, 0, -1):
                response, lines, octets = await loop.run_in_executor(None, server.retr, i)
                msg_content = b"\n".join(lines)
                msg = await loop.run_in_executor(None, BytesParser(policy=default).parsebytes, msg_content)

                # 获取发件人
                sender = msg.get('From')
                subject = msg.get('Subject')
                logger.info(f"邮件 {i} 发件人: {sender}")
                logger.info(f"邮件 {i} 主题: {subject}")

                # 检查是否是目标发件人
                if sender and target_sender.lower() in sender.lower():
                    logger.info(f"找到来自 {target_sender} 的邮件，正在处理...")

                    # 获取邮件正文
                    body = parse_email_body(msg)
                    if body:
                        # 提取验证码
                        code = extract_verification_code(body)
                        if code:
                            logger.info(f"找到验证码: {code}")

                            # 删除已处理的邮件
                            await loop.run_in_executor(None, server.dele, i)
                            logger.info(f"已删除邮件 {i}")

                            # 关闭服务器连接
                            await loop.run_in_executor(None, server.quit)
                            server = None  # 设置 server 为 None，防止 finally 块再次调用 quit

                            return code
                        else:
                            logger.info("未找到验证码，继续查找下一封邮件。")
                else:
                    logger.info(f"邮件 {i} 不是来自 {target_sender} 的邮件。")

            # 如果遍历完所有邮件都没有找到验证码
            logger.info(f"第 {attempt} 次尝试未找到验证码。")
            if attempt < retries:
                logger.info(f"等待 {sleep_time} 秒后重试...")
                await asyncio.sleep(sleep_time)
            else:
                logger.warning(f"已达到最大重试次数 ({retries})，未找到验证码。")
                return None

        except Exception as e:
            logger.error(f"第 {attempt} 次尝试失败，错误信息: {e}")
            if 'unknown user name or bad password.' in e:
                return None
            if attempt < retries:
                logger.info(f"等待 {sleep_time} 秒后重试...")
                await asyncio.sleep(sleep_time)
            else:
                logger.error(f"已达到最大重试次数 ({retries})，操作失败。")
                return None
        finally:
            if server:
                try:
                    await loop.run_in_executor(None, server.quit)
                except Exception as e:
                    logger.error(f"关闭 POP3 服务器连接时出错: {e}")

    return None

# 将 connect_imap_with_retry 修改为异步函数
async def connect_imap_async(email, client_id, refresh_token, target_sender=None, retries=6, sleep_time=2):
    """
    异步连接 IMAP 服务器，遍历 INBOX 和 Junk 文件夹，查找来自 target_sender 的邮件，
    提取验证码并删除邮件。包含重试机制。
    """
    loop = asyncio.get_event_loop()

    for attempt in range(1, retries + 1):
        mail = None
        try:
            logger.info(f"第 {attempt} 次尝试连接 IMAP 服务器。")

            # 获取 access_token
            access_token = await get_access_token(client_id, refresh_token)

            # 使用 SSL 连接到 IMAP 服务器
            mail = await loop.run_in_executor(None, imaplib.IMAP4_SSL, 'outlook.office365.com')

            # 使用 OAuth2 验证
            auth_string = generate_auth_string(email, access_token)
            await loop.run_in_executor(None, mail.authenticate, 'XOAUTH2', lambda x: auth_string.encode())

            folders = ["INBOX", "Junk"]  # 要处理的文件夹列表

            for folder in folders:
                try:
                    # 选择文件夹
                    status, _ = await loop.run_in_executor(None, mail.select, folder)
                    if status != 'OK':
                        logger.error(f"无法选择文件夹 {folder}。")
                        continue

                    # 搜索所有邮件
                    status, messages = await loop.run_in_executor(None, mail.search, None, 'ALL')
                    if status != 'OK':
                        logger.error(f"搜索文件夹 {folder} 中的邮件失败。")
                        continue

                    mail_ids = messages[0].split()
                    num_messages = len(mail_ids)
                    logger.info(f"文件夹 {folder} 中共有 {num_messages} 封邮件。")

                    # 从最新的邮件开始遍历
                    for mail_id in reversed(mail_ids):
                        # 获取邮件内容
                        status, msg_data = await loop.run_in_executor(None, mail.fetch, mail_id, '(RFC822)')
                        if status != 'OK':
                            logger.warning(f"无法获取邮件 ID {mail_id.decode()}。")
                            continue

                        msg_content = msg_data[0][1]
                        msg = await loop.run_in_executor(None, BytesParser(policy=default).parsebytes, msg_content)

                        # 获取发件人和主题
                        sender = msg.get('From')
                        subject = msg.get('Subject')
                        logger.info(f"文件夹 {folder} 中的邮件 ID {mail_id.decode()} 发件人: {sender}")
                        logger.info(f"文件夹 {folder} 中的邮件 ID {mail_id.decode()} 主题: {subject}")

                        # 检查是否是目标发件人
                        if target_sender and sender and target_sender.lower() in sender.lower():
                            logger.info(f"找到来自 {target_sender} 的邮件，正在处理...")

                            # 获取邮件正文
                            body = parse_email_body(msg)
                            if body:
                                # 提取验证码
                                code = extract_verification_code(body)
                                if code:
                                    logger.info(f"找到验证码: {code}")

                                    # 删除已处理的邮件
                                    await loop.run_in_executor(None, mail.store, mail_id, '+FLAGS', '\\Deleted')
                                    logger.info(f"已标记删除邮件 ID {mail_id.decode()}")

                                    # 立即删除标记为删除的邮件
                                    await loop.run_in_executor(None, mail.expunge)

                                    await loop.run_in_executor(None, mail.logout)
                                    mail = None  # 防止 finally 中重复调用 logout

                                    return code
                                else:
                                    logger.info("未找到验证码，继续查找下一封邮件。")
                        else:
                            logger.info(f"邮件 {mail_id.decode()} 不是来自 {target_sender} 的邮件。")
                except imaplib.IMAP4.error as e:
                    logger.error(f"处理文件夹 {folder} 时出错: {e}")
                    if 'User is authenticated but not connected' in e:
                        return None
                    continue

            # 如果所有文件夹遍历完毕仍未找到目标邮件
            logger.info(f"第 {attempt} 次尝试未找到来自 {target_sender} 的邮件。")
            if attempt < retries:
                logger.info(f"等待 {sleep_time} 秒后重试...")
                await asyncio.sleep(sleep_time)
            else:
                logger.warning(f"已达到最大重试次数 ({retries})，未找到来自 {target_sender} 的邮件。")
                return None

        except imaplib.IMAP4.error as e:
            logger.error(f"第 {attempt} 次尝试连接失败，错误信息: {e}")
            if 'unknown user name or bad password.' in e:
                return None
            if attempt < retries:
                logger.info(f"等待 {sleep_time} 秒后重试...")
                await asyncio.sleep(sleep_time)
            else:
                logger.error(f"已达到最大重试次数 ({retries})，操作失败。")
                return None
        except Exception as e:
            logger.error(f"第 {attempt} 次尝试发生异常，错误信息: {e}")
            if 'unknown user name or bad password.' in e:
                return None
            if attempt < retries:
                logger.info(f"等待 {sleep_time} 秒后重试...")
                await asyncio.sleep(sleep_time)
            else:
                logger.error(f"已达到最大重试次数 ({retries})，操作失败。")
                return None
        finally:
            if mail:
                try:
                    await loop.run_in_executor(None, mail.close)
                except:
                    pass
                try:
                    await loop.run_in_executor(None, mail.logout)
                except:
                    pass

    return None

async def connect_host_async(email, password, email_service, prot, target_sender, retries=4, sleep_time=2):
    loop = asyncio.get_running_loop()

    for attempt in range(1, retries + 1):
        server = None
        try:
            logger.info(f"第 {attempt} 次尝试连接 POP3 服务器并获取验证码。")
            # 连接服务器
            server = await loop.run_in_executor(None, poplib.POP3, email_service, prot)

            # 登录
            await loop.run_in_executor(None, server.user, email)
            await loop.run_in_executor(None, server.pass_, password)

            # mail = poplib.POP3(email_service, 110)
            logger.info(f"成功登录邮箱 {email}")

            # 获取邮件列表
            response, listings, octets = await loop.run_in_executor(None, server.list)
            num_messages = len(listings)
            logger.info(f"收件箱中共有 {num_messages} 封邮件。")

            # 从最新的邮件开始遍历
            for i in range(num_messages, 0, -1):
                response, lines, octets = await loop.run_in_executor(None, server.retr, i)
                msg_content = b"\n".join(lines)
                msg = await loop.run_in_executor(None, lambda: BytesParser(policy=default).parsebytes(msg_content))

                # 获取发件人
                sender = msg.get('From')
                subject = msg.get('Subject')
                logger.info(f"邮件 {i} 发件人: {sender}")
                logger.info(f"邮件 {i} 主题: {subject}")

                # 检查是否是目标发件人
                if sender and target_sender.lower() in sender.lower():
                    logger.info(f"找到来自 {target_sender} 的邮件，正在处理...")

                    # 获取邮件正文
                    body = await loop.run_in_executor(None, parse_email_body, msg)
                    if body:
                        # 提取验证码
                        code = await loop.run_in_executor(None, extract_verification_code, body)
                        if code:
                            logger.info(f"找到验证码: {code}")

                            # 删除已处理的邮件
                            await loop.run_in_executor(None, server.dele, i)
                            logger.info(f"已删除邮件 {i}")

                            # 关闭服务器连接
                            await loop.run_in_executor(None, server.quit)
                            server = None

                            return code
                        else:
                            logger.info("未找到验证码，继续查找下一封邮件。")
                else:
                    logger.info(f"邮件 {i} 不是来自 {target_sender} 的邮件。")

            logger.info(f"第 {attempt} 次尝试未找到验证码。")
            if attempt < retries:
                logger.info(f"等待 {sleep_time} 秒后重试...")
                await asyncio.sleep(sleep_time)
            else:
                logger.warning(f"已达到最大重试次数 ({retries})，未找到验证码。")
                return None

        except Exception as e:
            logger.error(f"第 {attempt} 次尝试失败，错误信息: {e}")
            if 'unknown user name or bad password.' in str(e):
                return None
            if attempt < retries:
                logger.info(f"等待 {sleep_time} 秒后重试...")
                await asyncio.sleep(sleep_time)
            else:
                logger.error(f"已达到最大重试次数 ({retries})，操作失败。")
                return None
        finally:
            if server:
                try:
                    await loop.run_in_executor(None, server.quit)
                except Exception as e:
                    logger.error(f"关闭 POP3 服务器连接时出错: {e}")

    return None